use std::{
    io,
    net::SocketAddr,
    os::fd::AsRawFd,
    pin::Pin,
    task::{Context, Poll},
};

use tokio::{
    io::{AsyncRead, AsyncWrite, DuplexStream, ReadBuf},
    net::{TcpStream, UnixStream},
};

use crate::utils::unique_id::UniqueID;

pub type MockStream = DuplexStream;

#[derive(Debug)]
enum RawStream {
    Tcp(TcpStream),
    Unix(UnixStream),
    Mock(MockStream),
}

impl AsyncRead for RawStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        // Safety: Basic enum pin projection
        unsafe {
            match &mut Pin::get_unchecked_mut(self) {
                RawStream::Tcp(tcp_stream) => {
                    Pin::new_unchecked(tcp_stream).poll_read(cx, buf)
                }
                RawStream::Unix(unix_stream) => {
                    Pin::new_unchecked(unix_stream).poll_read(cx, buf)
                }
                RawStream::Mock(mock_stream) => {
                    Pin::new_unchecked(mock_stream).poll_read(cx, buf)
                }
            }
        }
    }
}

impl AsyncWrite for RawStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        // Safety: Basic enum pin projection
        unsafe {
            match &mut Pin::get_unchecked_mut(self) {
                RawStream::Tcp(tcp_stream) => {
                    Pin::new_unchecked(tcp_stream).poll_write(cx, buf)
                }
                RawStream::Unix(unix_stream) => {
                    Pin::new_unchecked(unix_stream).poll_write(cx, buf)
                }
                RawStream::Mock(mock_stream) => {
                    Pin::new_unchecked(mock_stream).poll_write(cx, buf)
                }
            }
        }
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        // Safety: Basic enum pin projection
        unsafe {
            match &mut Pin::get_unchecked_mut(self) {
                RawStream::Tcp(tcp_stream) => {
                    Pin::new_unchecked(tcp_stream).poll_flush(cx)
                }
                RawStream::Unix(unix_stream) => {
                    Pin::new_unchecked(unix_stream).poll_flush(cx)
                }
                RawStream::Mock(mock_stream) => {
                    Pin::new_unchecked(mock_stream).poll_flush(cx)
                }
            }
        }
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), std::io::Error>> {
        // Safety: Basic enum pin projection
        unsafe {
            match &mut Pin::get_unchecked_mut(self) {
                RawStream::Tcp(tcp_stream) => {
                    Pin::new_unchecked(tcp_stream).poll_shutdown(cx)
                }
                RawStream::Unix(unix_stream) => {
                    Pin::new_unchecked(unix_stream).poll_shutdown(cx)
                }
                RawStream::Mock(mock_stream) => {
                    Pin::new_unchecked(mock_stream).poll_shutdown(cx)
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct Stream {
    raw_stream: RawStream,
    _unique_id: UniqueID,
    _send_bytes: usize,
    _recv_bytes: usize,
    _peer_addr_str: String,
    _local_addr_str: String,
}

impl From<TcpStream> for Stream {
    fn from(s: TcpStream) -> Self {
        let _peer_addr_str = match s.peer_addr() {
            Ok(addr) => addr.to_string(),
            Err(_) => "invalid_peer_addr".to_string(),
        };

        let _local_addr_str = match s.local_addr() {
            Ok(addr) => addr.to_string(),
            Err(_) => "invalid_local_addr".to_string(),
        };

        Stream {
            raw_stream: RawStream::Tcp(s),
            _unique_id: UniqueID::new(),
            _send_bytes: 0,
            _recv_bytes: 0,
            _peer_addr_str,
            _local_addr_str,
        }
    }
}

impl From<UnixStream> for Stream {
    fn from(s: UnixStream) -> Self {
        let _peer_addr_str = match s.peer_addr() {
            Ok(addr) => match addr.as_pathname() {
                Some(p) => {
                    p.to_str().unwrap_or("invalid_unix_path").to_string()
                }
                None => "invalid_unix_path".to_string(),
            },
            Err(_) => "invalid_unix_path".to_string(),
        };

        let _local_addr_str = match s.local_addr() {
            Ok(addr) => match addr.as_pathname() {
                Some(p) => {
                    p.to_str().unwrap_or("invalid_unix_path").to_string()
                }
                None => "invalid_unix_path".to_string(),
            },
            Err(_) => "invalid_unix_path".to_string(),
        };

        Stream {
            raw_stream: RawStream::Unix(s),
            _unique_id: UniqueID::new(),
            _send_bytes: 0,
            _recv_bytes: 0,
            _peer_addr_str,
            _local_addr_str,
        }
    }
}

impl From<MockStream> for Stream {
    fn from(s: MockStream) -> Self {
        Stream {
            raw_stream: RawStream::Mock(s),
            _unique_id: UniqueID::new(),
            _send_bytes: 0,
            _recv_bytes: 0,
            _peer_addr_str: "mock_stream".to_string(),
            _local_addr_str: "mock_stream".to_string(),
        }
    }
}

#[allow(dead_code)]
impl Stream {
    #[inline]
    pub fn unique_id(&self) -> UniqueID {
        self._unique_id
    }

    #[inline]
    pub fn as_raw_fd(&self) -> Option<std::os::unix::io::RawFd> {
        match &self.raw_stream {
            RawStream::Tcp(tcp_stream) => Some(tcp_stream.as_raw_fd()),
            RawStream::Unix(unix_stream) => Some(unix_stream.as_raw_fd()),
            RawStream::Mock(_) => None,
        }
    }

    #[inline]
    pub fn send_bytes(&self) -> usize {
        self._send_bytes
    }

    #[inline]
    pub fn recv_bytes(&self) -> usize {
        self._recv_bytes
    }

    #[inline]
    pub fn local_addr(&self) -> Option<SocketAddr> {
        match &self.raw_stream {
            RawStream::Tcp(tcp_stream) => tcp_stream.local_addr().ok(),
            RawStream::Unix(_) => None,
            RawStream::Mock(_) => None,
        }
    }

    #[inline]
    pub fn peer_addr(&self) -> Option<SocketAddr> {
        match &self.raw_stream {
            RawStream::Tcp(tcp_stream) => tcp_stream.peer_addr().ok(),
            RawStream::Unix(_) => None,
            RawStream::Mock(_) => None,
        }
    }

    #[inline]
    pub fn peer_addr_str(&self) -> &str {
        &self._peer_addr_str
    }

    #[inline]
    pub fn local_addr_str(&self) -> &str {
        &self._local_addr_str
    }
}

impl AsyncRead for Stream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let old_len = buf.filled().len();

        let res = Pin::new(&mut self.raw_stream).poll_read(cx, buf);

        let new_len = buf.filled().len();
        if new_len > old_len {
            let nr = new_len - old_len;
            self._recv_bytes += nr;
        }

        res
    }
}

impl AsyncWrite for Stream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        let res = Pin::new(&mut self.raw_stream).poll_write(cx, buf);
        if let Poll::Ready(Ok(n)) = res {
            self._send_bytes += n;
        }
        res
    }

    fn poll_flush(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.raw_stream).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.raw_stream).poll_shutdown(cx)
    }
}

// 创建Mock Stream, 返回ClientStream和ServerStream
#[allow(dead_code)]
pub fn create_mock_stream() -> (Stream, Stream) {
    let (client, server) = tokio::io::duplex(4096);
    (client.into(), server.into())
}

#[cfg(test)]
mod tests {

    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    use crate::{runtime::RuntimeManager, utils::fixed_buffer::FixedBuffer};

    use super::create_mock_stream;

    #[test]
    fn test_mock_stream() {
        RuntimeManager::init_main_runtime(1);
        RuntimeManager::main_runtime_block_on(async {
            let mut buf1 = FixedBuffer::new(1024);
            let (mut client_stream, mut server_stream) = create_mock_stream();
            client_stream.write(b"aaaa").await.unwrap();

            let write_slice = buf1.write_slice();
            let n = server_stream.read(write_slice).await.unwrap();
            buf1.fill(n);

            let read_slice = buf1.read_slice();
            assert!(read_slice.eq(b"aaaa"));

            let mut buf2 = FixedBuffer::new(1024);
            server_stream.write(b"bbbb").await.unwrap();
            let write_slice = buf2.write_slice();
            let n = client_stream.read(write_slice).await.unwrap();
            buf2.fill(n);

            let read_slice = buf2.read_slice();
            assert!(read_slice.eq(b"bbbb"));
        });
    }
}
